'''
@author: Hyunmin Kim, Euncheon Lim

JC69 model
JC69 is the simplest substitution model.
There are several assumptions.
It assumes equal base frequencies \left(\pi_A = \pi_G = \pi_C = \pi_T = {1\over4}\right) and equal mutation rates.
The only parameter of this model is therefore \mu, the overall substitution rate.
As previously mentioned, this variable becomes a constant when we normalize to the mean-rate to 1

'''
from abcframe import *
import numpy 
from scipy import linalg
import matplotlib.pyplot as plt
import string

class JC69Model(Model):
	_Q = None;
	_pi = [0.25,0.25,0.25,0.25];
	_probs = {}
	def __init__(self):
		Q = numpy.mat('[-3 1 1 1; 1 -3 1 1; 1 1 -3 1; 1 1 1 -3]')/4.0 
		s = 0.0
		for i in range(4):
			s -= Q[i,i] * self._pi[i]
		beta = 1/s		
		# normalize the Q independent of mu rates
		self._Q = beta * Q
		return
		
	## log likelihood calculation
	## exp(Q * t) =  P
	def prob(self, nu1, nu2, t):
		P = None
		## find precalculated probability 
		## Todo: granularization to reduce the hash size
		if self._probs.has_key(t):
			P = self._probs[t]
		else:
			P = linalg.expm(self._Q * t)
			self._probs[t]  = P
		return self._pi[nu1]*P[nu1,nu2]

	## this will be used when MCMC is involved 
	def update(self):
		return NotImplemented

if  __name__ =='__main__':
	model = JC69Model()
	print model.prob(1,2,0.21) 
	print model._Q
	print model._probs
	print model._probs[0.21]


